from torch import tensor, Tensor
from torch.utils.data import DataLoader
from torch.nn.utils.rnn import pad_sequence
from utils import conf, data
from .dataset import MovieDataset

_data, word2id = data.build_vocab()

def collate_fn(batch: list) -> tuple[Tensor, Tensor, Tensor]:
    # Convert word to id
    x_train = [tensor([word2id[x] for x in d[0]]) for d in batch]
    # Convert tag to id
    y_train = [tensor([data.labels[y] for y in d[1]]) for d in batch]
    # Pad sequence
    input_ids = pad_sequence(
        x_train,
        batch_first=True,
        padding_value=0
    )
    # Pad sequence
    labels = pad_sequence(
        y_train,
        batch_first=True,
        padding_value=data.labels['PAD']
    )
    # Create attention mask
    attention_mask = (input_ids != 0).long()
    # Return input ids, labels and attention mask
    return input_ids, labels, attention_mask

def get_dataloader() -> tuple[DataLoader, DataLoader]:
    # Create train dataset
    train_dataset = MovieDataset(_data[:300])
    # Create valid dataset
    valid_dataset = MovieDataset(_data[300:])
    # Create train dataloader
    train_dataloader = DataLoader(
        train_dataset,
        batch_size=conf.model.batch_size,
        collate_fn=collate_fn,
        drop_last=True
    )
    # Create valid dataloader
    valid_dataloader = DataLoader(
        valid_dataset,
        batch_size=conf.model.batch_size,
        collate_fn=collate_fn,
        drop_last=True
    )
    # Return train and valid dataloader
    return train_dataloader, valid_dataloader
